import datetime
import os
import pprint
import time
import threading
import torch as th
from types import SimpleNamespace as SN
from utils.logging import Logger
from utils.timehelper import time_left, time_str
from os.path import dirname, abspath

from learners import REGISTRY as le_REGISTRY
from runners import REGISTRY as r_REGISTRY
from controllers import REGISTRY as mac_REGISTRY
from components.episode_buffer import ReplayBuffer
from components.transforms import OneHot

import json
import numpy as np
from copy import deepcopy
from tensorboard.backend.event_processing import event_accumulator

def run(_run, _config, _log):

    # check args sanity
    _config = args_sanity_check(_config, _log)

    args = SN(**_config)
    args.device = "cuda" if args.use_cuda else "cpu"

    # setup loggers
    logger = Logger(_log)

    _log.info("Experiment Parameters:")
    experiment_params = pprint.pformat(_config,
                                       indent=4,
                                       width=1)
    _log.info("\n\n" + experiment_params + "\n")

    # configure tensorboard logger
    results_save_dir = args.results_save_dir

    if args.use_tensorboard and not args.evaluate:
        # only log tensorboard when in training mode
        tb_exp_direc = os.path.join(results_save_dir, 'tb_logs')
        logger.setup_tb(tb_exp_direc)

        config_str = json.dumps(vars(args), indent=4)
        with open(os.path.join(results_save_dir, "config.json"), "w") as f:
            f.write(config_str)

    # sacred is on by default
    logger.setup_sacred(_run)

    # Run and train
    run_sequential(args=args, logger=logger)

    # Clean up after finishing
    print("Exiting Main")

    print("Stopping all threads")
    for t in threading.enumerate():
        if t.name != "MainThread":
            print("Thread {} is alive! Is daemon: {}".format(t.name, t.daemon))
            t.join(timeout=1)
            print("Thread joined")

    print("Exiting script")

    # Making sure framework really exits
    os._exit(os.EX_OK)

def run_sequential(args, logger):

    # Init runner so we can get env info
    runner = r_REGISTRY[args.runner](args=args, logger=logger)

    # Set up schemes and groups here
    env_info = runner.get_env_info()
    args.n_agents = env_info["n_agents"]
    args.n_actions = env_info["n_actions"]
    args.state_shape = env_info["state_shape"]

    # Default/Base scheme
    scheme = {
        "state": {"vshape": env_info["state_shape"]},
        "obs": {"vshape": env_info["obs_shape"], "group": "agents"},
        "actions": {"vshape": (1,), "group": "agents", "dtype": th.long},
        "avail_actions": {"vshape": (env_info["n_actions"],), "group": "agents", "dtype": th.int},
        "reward": {"vshape": (1,)},
        "terminated": {"vshape": (1,), "dtype": th.uint8},
    }
    groups = {
        "agents": args.n_agents
    }
    preprocess = {
        "actions": ("actions_onehot", [OneHot(out_dim=args.n_actions)])
    }

    runner.setup(scheme=scheme, groups=groups, preprocess=preprocess)
    empty_buffer = ReplayBuffer(scheme, groups, args.buffer_size, env_info["episode_limit"] + 1,
                          preprocess=preprocess,
                          device="cpu" if args.buffer_cpu_only else args.device)

    ## Setup ego
    mac_ego = mac_REGISTRY['multi_head_mac'](empty_buffer.scheme, groups, args)
    learner_ego = le_REGISTRY['multi_head_learner'](mac_ego, empty_buffer.scheme, logger, args)
    if args.use_cuda:
        learner_ego.cuda()

    def get_timestep(path, level):
        timesteps = []
        if not os.path.isdir(path):
            logger.console_logger.info("Checkpoint directiory {} doesn't exist".format(args.checkpoint_path))
            return
        # Go through all files in args.checkpoint_path
        for name in os.listdir(path):
            full_name = os.path.join(path, name)
            # Check if they are dirs the names of which are numbers
            if os.path.isdir(full_name) and name.isdigit():
                timesteps.append(int(name))

        if level == 2:
            return min(timesteps)
        elif level == 0:
            return max(timesteps)
        elif level == 1:
            tb_path = path.split('/')
            if '2024' in path:
                year = '2024'
            elif '2025' in path:
                year = '2025'
            elif '2026' in path:
                year = '2026'
            else:
                assert 0
            seed_name = '_'.join(tb_path[-1].split(f'_{year}')[0].split('_')[:2]) + '_'
            run_name = '_'.join(tb_path[-1].split(f'_{year}')[0].split('_')[2:])
            tb_path = tb_path[:-2]
            tb_path.append('pymarl')
            tb_path.append(args.env)
            tb_path.append(run_name)
            tb_path = '/'.join(tb_path)
            for _name in os.listdir(tb_path):
                if seed_name in _name:
                    tb_path = os.path.join(tb_path, _name, 'tb_logs')
                    break
            tb_path = os.path.join(tb_path, os.listdir(tb_path)[0])
            ea = event_accumulator.EventAccumulator(tb_path)
            ea.Reload()
            scalar_data = ea.scalars.Items('test_return_original_mean')
            
            try:
                steps = [scalar_data[i][1] for i in range(len(scalar_data))]
                values = [scalar_data[i][2] for i in range(len(scalar_data))]
            except:
                steps = [scalar_data[i].step for i in range(len(scalar_data))]
                values = [scalar_data[i].value for i in range(len(scalar_data))]
            mid_value = (values[-1] - values[0]) / 2
            min_delta = 1e9
            load_step = min(timesteps)
            for step, value in zip(steps, values):
                if min_delta > abs(mid_value - value):
                    min_delta = abs(mid_value - value)
                    load_step = step
            return min(timesteps, key=lambda x: abs(x - load_step))
        else:
            assert 0

    if args.last_ego_model_path != "":
        timestep_to_load_tm = get_timestep(args.last_ego_model_path, 0)
        model_path_ego = os.path.join(args.last_ego_model_path, str(timestep_to_load_tm))
        logger.console_logger.info(f"Loading ego model from {model_path_ego}")
        learner_ego.load_models(model_path_ego)
        mac_ego.agent.reset_head()
        # load previous head_dict, and create a new head
        learner_ego.cache_feature_layer()

    ## Setup tm
    args.checkpoint_path_tm_lst = args.checkpoint_path_tm_lst[1:-1].split(',')
    assert type(args.checkpoint_path_tm_lst) == type([]) and len(args.checkpoint_path_tm_lst) > 0
    n_checkpoints = 3 if args.utilize_checkpoints else 1
    args.n_population = len(args.checkpoint_path_tm_lst) * n_checkpoints

    if args.last_ego_model_path != "":
        mac_ego.agent.store_head(args.n_population // n_checkpoints - 1) # head expansion, watch out
    if args.last_ego_model_path == "": # LFS
        for i in range(args.n_population // n_checkpoints):
            mac_ego.agent.reset_head()
            mac_ego.agent.store_head(i)
    learner_ego.reset_params_and_optimizer()
    learner_ego._update_targets()

    mac_tm_init = mac_REGISTRY[args.mac](empty_buffer.scheme, groups, args)
    tm2learner, tm2mac = {}, {}
    for tm_id_init in range(args.n_population):
        mac_tm = deepcopy(mac_tm_init)
        tm2learner[tm_id_init] = le_REGISTRY[args.learner](mac_tm, empty_buffer.scheme, logger, args)
        if args.use_cuda:
            tm2learner[tm_id_init].cuda()
        tm2mac[tm_id_init] = mac_tm
    
    for load_tm_id_, checkpoint_path_tm in enumerate(args.checkpoint_path_tm_lst):
        for i in range(n_checkpoints):
            load_tm_id = load_tm_id_ * n_checkpoints + i
            timestep_to_load_tm = get_timestep(checkpoint_path_tm, i)
            model_path_tm = os.path.join(checkpoint_path_tm, str(timestep_to_load_tm))
            logger.console_logger.info(f"Loading tm {load_tm_id} model from {model_path_tm}")
            tm2learner[load_tm_id].load_models(model_path_tm)
            
    tm2buffer = {tm_idx: deepcopy(empty_buffer) for tm_idx in range(args.n_population // n_checkpoints)}

    # start training
    episode = 0
    last_test_T = -args.test_interval - 1
    last_log_T = 0
    model_save_time = 0

    start_time = time.time()
    last_time = start_time

    n_test_runs = max(1, args.test_nepisode // runner.batch_size)
    for tm_test, mac_test in tm2mac.items():
        for _ in range(n_test_runs):
            runner.run(mac1=mac_ego,  mac2=mac_test, test_mode=True, test_mode_1=True, test_mode_2=True, tm_id=tm_test, head_id=int(tm_test / n_checkpoints))

    # args.t_max *= (args.n_population // n_checkpoints)
    logger.console_logger.info("Beginning training for {} timesteps".format(args.t_max))

    while runner.t_env <= args.t_max:

        # tm_id_train = np.random.randint(args.n_population)
        tm_id_train = np.random.choice(np.arange(args.n_population)[-1 * n_checkpoints:]) # only train with the latest tm

        ## collect xp traj and learn
        episode_batch = runner.run(mac1=mac_ego, mac2=tm2mac[tm_id_train], test_mode=False, test_mode_1=False, test_mode_2=True, tm_id=tm_id_train, head_id=int(tm_id_train / n_checkpoints), eps_greedy_t=runner.t_env)
        tm2buffer[int(tm_id_train / n_checkpoints)].insert_episode_batch(episode_batch)
        if tm2buffer[int(tm_id_train / n_checkpoints)].can_sample(args.batch_size):
            episode_sample = tm2buffer[int(tm_id_train / n_checkpoints)].sample(args.batch_size)
            max_ep_t = episode_sample.max_t_filled()
            episode_sample = episode_sample[:, :max_ep_t]
            if episode_sample.device != args.device:
                episode_sample.to(args.device)
            learner_ego.train(episode_sample, runner.t_env, episode, int(tm_id_train / n_checkpoints))

        # Execute test runs once in a while
        n_test_runs = max(1, args.test_nepisode // runner.batch_size)
        if (runner.t_env - last_test_T) / args.test_interval >= 1.0:

            logger.console_logger.info("t_env: {} / {}".format(runner.t_env, args.t_max))
            logger.console_logger.info("Estimated time left: {}. Time passed: {}".format(
                time_left(last_time, last_test_T, runner.t_env, args.t_max), time_str(time.time() - start_time)))
            last_time = time.time()

            last_test_T = runner.t_env
            for tm_test, mac_test in tm2mac.items():
                for _ in range(n_test_runs):
                    runner.run(mac1=mac_ego, mac2=mac_test, test_mode=True, test_mode_1=True, test_mode_2=True, tm_id=tm_test, head_id=int(tm_test / n_checkpoints))

        if args.save_model and (runner.t_env - model_save_time >= args.save_model_interval or model_save_time == 0):
            model_save_time = runner.t_env
            save_path = os.path.join(args.local_results_path, "models", args.unique_token, str(runner.t_env))
            #"results/models/{}".format(unique_token)
            os.makedirs(save_path, exist_ok=True)
            logger.console_logger.info("Saving models to {}".format(save_path))

            # learner should handle saving/loading -- delegate actor save/load to mac,
            # use appropriate filenames to do critics, optimizer states
            learner_ego.save_models(save_path)

        episode += args.batch_size_run

        if (runner.t_env - last_log_T) >= args.log_interval:
            logger.log_stat("episode", episode, runner.t_env)
            logger.print_recent_stats()
            last_log_T = runner.t_env

    if args.save_model:
        save_path = os.path.join(args.local_results_path, "models", args.unique_token, str(runner.t_env))
        #"results/models/{}".format(unique_token)
        os.makedirs(save_path, exist_ok=True)
        logger.console_logger.info("Saving models to {}".format(save_path))

        # learner should handle saving/loading -- delegate actor save/load to mac,
        # use appropriate filenames to do critics, optimizer states
        learner_ego.save_models(save_path)

    runner.close_env()
    logger.console_logger.info("Finished Training")

def args_sanity_check(config, _log):

    # set CUDA flags
    # config["use_cuda"] = True # Use cuda whenever possible!
    if config["use_cuda"] and not th.cuda.is_available():
        config["use_cuda"] = False
        _log.warning("CUDA flag use_cuda was switched OFF automatically because no CUDA devices are available!")

    if config["test_nepisode"] < config["batch_size_run"]:
        config["test_nepisode"] = config["batch_size_run"]
    else:
        config["test_nepisode"] = (config["test_nepisode"]//config["batch_size_run"]) * config["batch_size_run"]

    return config